package cn.itcast.manus.agent.planning;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.itcast.manus.agent.Agent;
import cn.itcast.manus.agent.AgentFactory;
import cn.itcast.manus.agent.ReActBaseAgent;
import cn.itcast.manus.agent.prompt.PromptManagement;
import cn.itcast.manus.config.ModelConfig;
import cn.itcast.manus.constants.Constant;
import cn.itcast.manus.enums.AgentTypeEnum;
import cn.itcast.manus.message.MessageSession;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Function;

/**
 * ReActPlanningAgent 是一个任务规划代理类，继承自 ReActBaseAgent，
 * 专门负责将复杂任务分解为可执行的子任务序列，为后续的具体执行代理提供行动计划。
 * 它是整个 ReAct 代理架构中的规划层组件。
 * 假设有以下输入：
 * 任务：分析近3年全球芯片产业发展趋势，并制作相关数据可视化图表
 * ReActPlanningAgent 可能生成的规划结果：
 * {
 *   "status": {
 *     "thinking": "我需要先收集近3年全球芯片产业的数据，然后分析发展趋势，最后制作可视化图表。",
 *     "memory": "",
 *     "evaluationPreviousGoal": ""
 *   },
 *   "action": [
 *     {
 *       "search": {
 *         "query": "2021-2024年全球芯片产业市场规模数据",
 *         "maxResults": 5
 *       }
 *     }
 *   ]
 * }
 * 通过这种方式，ReActPlanningAgent 实现了将高层任务转换为具体可执行步骤的关键功能，在复杂任务处理系统中扮演着"大脑"的角色。
 */
@Slf4j
public class ReActPlanningAgent extends ReActBaseAgent {

    @Resource
    private PromptManagement promptManagement;
    @Resource(name = ModelConfig.PLAN_AGENT)
    private ChatModel chatModel;
    // 存放agent，key：agentName，value：功能描述
    private final Map<String, String> agentInfo = new HashMap<>();
    // 线程安全的动态数组，用来存储子任务链
    private final CopyOnWriteArrayList<SubTaskNode> subTaskChain = new CopyOnWriteArrayList<>();

    public ReActPlanningAgent(MessageSession messageSession) {
        super(messageSession);
    }
    @PostConstruct
    void init() {
        // 将需要的智能体名称和描述添加到智能体信息中
        agentInfo.put(AgentTypeEnum.BROWSER_AGENT.getAgentName(), AgentTypeEnum.BROWSER_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.TABLE_AGENT.getAgentName(), AgentTypeEnum.TABLE_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.MARKDOWN_AGENT.getAgentName(), AgentTypeEnum.MARKDOWN_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.CHART_AGENT.getAgentName(), AgentTypeEnum.CHART_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.HTML_DOC_AGENT.getAgentName(), AgentTypeEnum.HTML_DOC_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.AMAP_AGENT.getAgentName(), AgentTypeEnum.AMAP_AGENT.getDesc());
        agentInfo.put(AgentTypeEnum.DB_AGENT.getAgentName(), AgentTypeEnum.DB_AGENT.getDesc());

    }

    // 目标任务
    private String targetTask;
    // 子任务中间结果
    private final Map<String, String> subTaskMidResult = new ConcurrentHashMap<>();

    /**
     * 这个设计体现了继承与组合的设计模式思想，
     * ReActPlanningAgent 专注于规划特有的逻辑（子任务链管理），
     * 而将通用的 ReAct 推理流程委托给父类处理，实现了代码的有效复用和功能分离。
     * @param task
     * @return
     */
    @Override
    public String reActSolve(String task) {
        this.subTaskChain.clear(); //清空子任务链，确保每次处理新任务时都从一个干净的状态开始
        this.subTaskMidResult.clear();
        this.targetTask = task;
        return super.reActSolve(task); //然后调用父类 ReActBaseAgent 的 reActSolve(task) 方法执行实际的 ReAct 推理流程
        /*  reActSolve
            初始化子任务链
            初始化步骤计数器和动作映射
            构建包含任务信息的输入消息列表
            循环执行思考-行动-观察过程，直到任务完成或达到最大步数限制
            调用大模型处理每个步骤的决策
         */
    }

    @Override
    protected String getCurrentStatus() {
        // 获取任务链中的最后一个任务，查找到对应的Agent进行执行
        // 获取任务链中的最后一个任务，查找到对应的Agent进行执行
        this.subTaskChain.stream()
                .reduce((x, y) -> y)
                .ifPresent(last -> {
                    // 获取当前任务
                    String originalTask = last.getSubTask();
                    // 将原任务与目标任务合并
                    String subTaskWithCtx = this.taskMerging(targetTask, originalTask, subTaskMidResult);

                    // 查找对应的智能体，修改最大步数为当前任务的最大步数
                    Function<MessageSession, Agent> agentFunction = AgentFactory.getAgent(AgentTypeEnum.agentNameOf(last.getAgent()));
                    if (ObjectUtil.isEmpty(agentFunction)) {
                        return;
                    }
                    var agent = agentFunction.apply(super.messageSession);
                    if (agent instanceof ReActBaseAgent base) {
                        // 修改任务的最大步数为最后一个任务的最大步数
                        base.resetReActConfig(c -> c.setMaxStep(last.getMaxStep()));
                    }

                    // 调用智能体，获取结果
                    String result = agent.solveTask(subTaskWithCtx);
                    if (agent instanceof ReActBaseAgent base) {
                        // 设置结果的步数，就是当前的任务的步数
                        last.setResultStep(base.getCurrentStep());
                    }
                    // 将结果保存到中间结果中
                    subTaskMidResult.put(originalTask, result);
                    last.setResult(result);
                });
        // 将子任务链中的结果以字符串拼接，拼接格式为：子任务1->子任务2->子任务3
        var chainStr = subTaskChain.stream()
                .map(SubTaskNode::strInStatus)
                .reduce((x, y) -> x + "->" + y)
                .orElse("[Empty]");

        // 将子任务链中最后一个任务的结果拼接
        var latestStr = subTaskChain.stream()
                .reduce((x, y) -> y)
                .map(SubTaskNode::getResult)
                .orElse("[Empty]");

        // 将结果拼接到模板中，作为下次调用大模型的历史对话
        var params = Map.of(
                "stepData", StrUtil.format("{}/{}", super.currentStep, super.reActConfig.getMaxStep()),
                "dateTime", DateUtil.now(),
                "subTaskChain", chainStr,
                "latestResult", latestStr);
        return StrUtil.format(this.promptManagement.getPrompt(Constant.Prompts.PLANNING_STATUS), params);
    }
    private String taskMerging(String targetTask, String subTask, Object context) {
        String prompt = this.promptManagement.getPrompt(Constant.Prompts.PLANNING_TASK_MERGING);
        var params = Map.of(
                "targetTask", targetTask,
                "subTask", subTask,
                "context", Convert.toStr(context));
        return StrUtil.format(prompt, params);
    }
    @Override
    protected boolean isStatusSignificantChanged() {
        // 计划Agent每次的action只需要执行一个，后面的action不再执行
        return true;
    }

    /**
     * 该类通过重写 addInitInput 方法，为大模型构建特定的规划上下文：
     * @param inputList
     * @param task
     */
    @Override
    protected void addInitInput(List<Message> inputList, String task) {
        // 添加系统消息 包含规划代理的角色定义和工作流程
        inputList.add(SystemMessage.builder()
                .text(promptManagement.getPrompt(Constant.Prompts.PLANNING_SYSTEM))
                .build());

        // 添加用户消息 包含具体任务和可用agent信息
        inputList.add(UserMessage.builder()
                .text(this.getPromptPlanningUserTask(task))
                .build());

        // 添加说明  添加记忆分割标记
        inputList.add(UserMessage.builder()
                .text("[你的历史记忆从此处开始]")
                .build());
    }

    /*
     * 方法将原始任务和可用代理信息格式化为大模型可理解的提示：
     */
    private String getPromptPlanningUserTask(String task) {
        var promptPlanningSystem = this.promptManagement.getPrompt(Constant.Prompts.PLANNING_USER_TASK);
        var param = Map.of(
                Constant.TASK, task, // 原始用户任务
                Constant.AGENT_DATA, String.valueOf(agentInfo)); // 可用代理信息
        return StrUtil.format(promptPlanningSystem, param);
    }
    /**
     * 规划下一个子任务节点
     * 可以把 ReActPlanningAgent 想象成一个项目经理：
     *
     * 接到一个大项目（复杂任务）
     * 清空办公桌（清空子任务链）
     * 将项目分解成多个小任务（规划子任务）
     * 为每个小任务分配给最合适的团队成员（指定agent）
     * 规定每个任务的完成时间（maxStep）
     * 记录每个任务的完成情况（result）
     * 根据前面任务的完成情况，继续规划下一个任务
     * 通过这种方式，即使是非常复杂的项目，也能有条不紊地推进和完成。
     *
     * @param agent  执行该子任务的agent名称
     * @param subTask  子任务内容
     * @param maxStep  完成任务的最大步数
     * @return  规划结果
     */
    @Tool(description = "规划下一个子任务节点")
    public String generateNext(@ToolParam(description = "agent名称") String agent,
                               @ToolParam(description = "子任务内容") String subTask,
                               @ToolParam(description = "完成任务的最大步数") int maxStep) {
        //通过流式 API 检查当前子任务链中是否有已执行的任务，据此生成不同的规划描述
        var res = subTaskChain.stream()
                .reduce((x, y) -> y) // 获取最后一个元素
                .map(p -> { // 对最后一个任务进行描述，生成了新的子任务
                    var result = "子任务{}的执行结果:\n{}\n基于此结果成功规划下一个子任务节点:{}";
                    return StrUtil.format(result, p.getSubTask(), p.getResult(), subTask);
                })
                // 没有子任务，就生成新的任务
                .orElseGet(() -> StrUtil.format("成功规划下一个子任务节点:{}", subTask));

        // 生成新的子任务节点
        SubTaskNode subTaskNode = SubTaskNode.builder()
                .agent(agent)
                .subTask(subTask)
                .maxStep(maxStep)
                .build();
        // 将子任务加入到任务链中
        this.subTaskChain.add(subTaskNode);
        return res;
    }

    @Override
    protected List<ToolCallbackProvider> toolCallbackProvider() {
        // 将当前对象中标记为@Tool解析出来，作为工具回调
        return List.of(() -> ToolCallbacks.from(ReActPlanningAgent.this));
    }


    /**
     * SubTaskNode 类是一个数据结构，用于表示任务分解过程中的单个子任务节点：
     *
     * 设计特点：使用了 Lombok 注解简化代码，包括 @Data（生成 getter/setter/toString 等）、@Builder（构建者模式）、@NoArgsConstructor 和 @AllArgsConstructor
     * 核心属性：包含执行子任务的智能体名称、子任务内容、最大步数限制、执行结果和结果对应的步骤等信息
     * 辅助方法：strInStatus() 方法提供了子任务节点的简洁字符串表示，用于日志记录和状态跟踪
     */
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class SubTaskNode {
        private String agent;           // 执行子任务的智能体名称
        private String subTask;         // 子任务内容
        private int maxStep;            // 完成子任务的最大步数
        private String result;          // 子任务执行结果
        private int resultStep = -1;    // 结果对应的步骤

        public String strInStatus() {
            return StrUtil.format("[{},Task:{},Result:{},resultStep:{}]",
                    agent,
                    subTask,
                    StrUtil.truncateUtf8(result, 20),
                    resultStep);
        }
    }

    @Override
    public ChatModel chatModel() {
        return this.chatModel;
    }
}
